#include "KDTree.h"
#include <memory>
#include <vector>
#include <algorithm>
#include <thread>
#include <cmath>
#include <iostream>

using namespace std;

// KDNode 成员函数实现
KDTree::KDNode::KDNode(const FlightLine& ln, int ax, double sv)
    : line(ln), axis(ax), splitValue(sv), bbox(ln.getBoundingBox()) {}

void KDTree::KDNode::updateBoundingBox() {
    bbox = line.getBoundingBox();
    if (left) mergeBBox(left->bbox);
    if (right) mergeBBox(right->bbox);
    if (auto p = parent.lock()) p->updateBoundingBox();
}

void KDTree::KDNode::mergeBBox(const pair<pair<double, double>, pair<double, double>>& b) {
    bbox.first.first = min(bbox.first.first, b.first.first);
    bbox.first.second = min(bbox.first.second, b.first.second);
    bbox.second.first = max(bbox.second.first, b.second.first);
    bbox.second.second = max(bbox.second.second, b.second.second);
}

void KDTree::KDNode::insert(const FlightLine& newLine, int depth) {
    int currentAxis = depth % 2;
    double newMid = currentAxis == 0 ? (newLine.x1 + newLine.x2) / 2 : (newLine.y1 + newLine.y2) / 2;

    if (newMid < splitValue) {
        if (!left) {
            left = make_shared<KDNode>(newLine, currentAxis, newMid);
            left->parent = shared_from_this();
            updateBoundingBox();
        }
        else left->insert(newLine, depth + 1);
    }
    else {
        if (!right) {
            right = make_shared<KDNode>(newLine, currentAxis, newMid);
            right->parent = shared_from_this();
            updateBoundingBox();
        }
        else right->insert(newLine, depth + 1);
    }
}

bool KDTree::KDNode::remove(int targetId, shared_ptr<KDNode>& root) {
    if (line.id == targetId) {
        // Case 1: 叶子节点
        if (!left && !right) {
            if (auto p = parent.lock()) {
                // 父节点断开链接
                if (p->left.get() == this)
                    p->left = nullptr;
                else
                    p->right = nullptr;
                p->updateBoundingBox();
            }
            else {
                root = nullptr; // 根节点置空
            }
            return true;
        }

        // Case 2: 只有一个子树
        if (!left || !right) {
            auto child = left ? left : right;

            if (auto p = parent.lock()) {
                // 父节点重新链接到孙子节点
                if (p->left.get() == this)
                    p->left = child;
                else
                    p->right = child;

                // 更新孙子节点的父指针
                if (child)
                    child->parent = p;

                p->updateBoundingBox();
            }
            else {
                // 更新根节点
                root = child;
                if (child)
                    child->parent.reset(); // 清除新根节点的父指针
            }
            return true;
        }

        // Case 3: 有两个子树
        // 找右子树中当前轴的最小节点
        auto minNode = findMinNode(right, axis);

        // 替换当前节点数据
        line = minNode->line;

        // 递归删除被替换的节点
        return right->remove(minNode->line.id, root);
    }

    // 递归搜索左右子树
    bool removed = false;
    if (left)
        removed = left->remove(targetId, root);

    if (!removed && right)
        removed = right->remove(targetId, root);

    // 更新当前节点的包围盒
    if (removed)
        updateBoundingBox();

    return removed;
}

shared_ptr<KDTree::KDNode> KDTree::KDNode::findMinNode(shared_ptr<KDNode> node, int targetAxis) {
    if (!node) return nullptr;

    // 比较当前节点
    auto best = node;

    // 比较左子树
    if (auto leftBest = findMinNode(node->left, targetAxis)) {
        double bestValue = (targetAxis == 0) ?
            (best->line.x1 + best->line.x2) / 2 :
            (best->line.y1 + best->line.y2) / 2;

        double leftValue = (targetAxis == 0) ?
            (leftBest->line.x1 + leftBest->line.x2) / 2 :
            (leftBest->line.y1 + leftBest->line.y2) / 2;

        if (leftValue < bestValue)
            best = leftBest;
    }

    // 比较右子树
    if (auto rightBest = findMinNode(node->right, targetAxis)) {
        double bestValue = (targetAxis == 0) ?
            (best->line.x1 + best->line.x2) / 2 :
            (best->line.y1 + best->line.y2) / 2;

        double rightValue = (targetAxis == 0) ?
            (rightBest->line.x1 + rightBest->line.x2) / 2 :
            (rightBest->line.y1 + rightBest->line.y2) / 2;

        if (rightValue < bestValue)
            best = rightBest;
    }

    return best;
}

shared_ptr<KDTree::KDNode> KDTree::KDNode::findMinNode(int targetAxis) {
    shared_ptr<KDNode> candidate = shared_from_this();

    if (left) {
        auto leftCandidate = left->findMinNode(targetAxis);
        double candidateValue = (targetAxis == 0) ?
            (candidate->line.x1 + candidate->line.x2) / 2 :
            (candidate->line.y1 + candidate->line.y2) / 2;
        double leftValue = (targetAxis == 0) ?
            (leftCandidate->line.x1 + leftCandidate->line.x2) / 2 :
            (leftCandidate->line.y1 + leftCandidate->line.y2) / 2;
        if (leftValue < candidateValue)
            candidate = leftCandidate;
    }

    if (right) {
        auto rightCandidate = right->findMinNode(targetAxis);
        double candidateValue = (targetAxis == 0) ?
            (candidate->line.x1 + candidate->line.x2) / 2 :
            (candidate->line.y1 + candidate->line.y2) / 2;
        double rightValue = (targetAxis == 0) ?
            (rightCandidate->line.x1 + rightCandidate->line.x2) / 2 :
            (rightCandidate->line.y1 + rightCandidate->line.y2) / 2;
        if (rightValue < candidateValue)
            candidate = rightCandidate;
    }

    return candidate;
}

void KDTree::KDNode::replaceChild(const shared_ptr<KDNode>& oldChild, const shared_ptr<KDNode>& newChild) {
    if (left == oldChild) left = newChild;
    else if (right == oldChild) right = newChild;
    if (newChild) newChild->parent = shared_from_this();
    updateBoundingBox();
}

// KDTree 成员函数实现
KDTree::KDTree(const vector<FlightLine>& lines) {
    vector<FlightLine> copy(lines);
    root = buildKDTree(copy);
}

void KDTree::InputFlightLine(const vector<FlightLine>& lines) {
    vector<FlightLine> copy(lines);
    root = buildKDTree(copy);
}

shared_ptr<KDTree::KDNode> KDTree::buildKDTree(vector<FlightLine>& lines, int depth) {
    if (lines.empty()) return nullptr;
    int axis = depth % 2;
    nth_element(lines.begin(), lines.begin() + lines.size() / 2, lines.end(),
        [axis](const FlightLine& a, const FlightLine& b) {
            double av = axis == 0 ? (a.x1 + a.x2) / 2 : (a.y1 + a.y2) / 2;
            double bv = axis == 0 ? (b.x1 + b.x2) / 2 : (b.y1 + b.y2) / 2;
            return av < bv;
        });
    auto node = make_shared<KDNode>(lines[lines.size() / 2], axis,
        axis == 0 ? (lines[lines.size() / 2].x1 + lines[lines.size() / 2].x2) / 2 :
        (lines[lines.size() / 2].y1 + lines[lines.size() / 2].y2) / 2);
    vector<FlightLine> leftLines(lines.begin(), lines.begin() + lines.size() / 2);
    vector<FlightLine> rightLines(lines.begin() + lines.size() / 2 + 1, lines.end());
    node->left = buildKDTree(leftLines, depth + 1);
    node->right = buildKDTree(rightLines, depth + 1);
    if (node->left) node->left->parent = node;
    if (node->right) node->right->parent = node;
    node->updateBoundingBox();
    return node;
}

double KDTree::pointToLineSq(const QueryPoint& p, const FlightLine& line) {
    double x1 = line.x1, y1 = line.y1, x2 = line.x2, y2 = line.y2;
    double dx = x2 - x1, dy = y2 - y1;
    if (dx == 0 && dy == 0) return (p.x - x1) * (p.x - x1) + (p.y - y1) * (p.y - y1);
    double t = max(0.0, min(1.0, ((p.x - x1) * dx + (p.y - y1) * dy) / (dx * dx + dy * dy)));
    double projX = x1 + t * dx, projY = y1 + t * dy;
    return (p.x - projX) * (p.x - projX) + (p.y - projY) * (p.y - projY);
}

double KDTree::pointToBBoxSq(const QueryPoint& p, const pair<pair<double, double>, pair<double, double>>& bbox) {
    double dx = max({ bbox.first.first - p.x, 0.0, p.x - bbox.second.first });
    double dy = max({ bbox.first.second - p.y, 0.0, p.y - bbox.second.second });
    return dx * dx + dy * dy;
}

void KDTree::nearestSearch(const shared_ptr<KDNode>& node, const QueryPoint& p, double& minSq, int& nearestId) const {
    if (!node) return;
    double bboxDist = pointToBBoxSq(p, node->bbox);
    if (bboxDist > minSq) return;

    double currDist = pointToLineSq(p, node->line);
    if (currDist < minSq || (currDist == minSq && node->line.id < nearestId)) {
        minSq = currDist;
        nearestId = node->line.id;
    }

    bool goLeft = (node->axis == 0 ? p.x : p.y) < node->splitValue;
    if (goLeft) {
        nearestSearch(node->left, p, minSq, nearestId);
        if (node->right) {
            double currentMinDist = sqrt(minSq);
            double diff = (node->axis == 0 ? p.x : p.y) - node->splitValue;
            if (diff <= currentMinDist) {
                nearestSearch(node->right, p, minSq, nearestId);
            }
        }
    }
    else {
        nearestSearch(node->right, p, minSq, nearestId);
        if (node->left) {
            double currentMinDist = sqrt(minSq);
            double diff = node->splitValue - (node->axis == 0 ? p.x : p.y);
            if (diff <= currentMinDist) {
                nearestSearch(node->left, p, minSq, nearestId);
            }
        }
    }
}

void KDTree::insert(const FlightLine& line) {
    if (!root) {
        int axis = 0;
        double splitValue = (line.x1 + line.x2) / 2;
        root = make_shared<KDNode>(line, axis, splitValue);
    }
    else {
        root->insert(line);
    }
}

bool KDTree::remove(int targetId) {
    if (!root) return false;
    return root->remove(targetId, root);
}

vector<KDTree::QueryResult> KDTree::findNearestBatch(const vector<QueryPoint>& queries) const {
    vector<QueryResult> results(queries.size(), { -1, numeric_limits<double>::quiet_NaN() });

    if (!root) return results;

    const size_t numQueries = queries.size();

    // 小数据集单线程处理
    if (numQueries < 200) {
        for (size_t i = 0; i < numQueries; ++i) {
            double minSq = numeric_limits<double>::max();
            int nearestId = -1;
            nearestSearch(root, queries[i], minSq, nearestId);
            double distance = (minSq == numeric_limits<double>::max()) ?
                numeric_limits<double>::quiet_NaN() : sqrt(minSq);
            results[i] = { nearestId, distance };
        }
        return results;
    }

    // 获取线程数并分配任务
    const size_t numThreads = thread::hardware_concurrency();
    vector<thread> threads;
    vector<vector<size_t>> threadIndices(numThreads);

    // 轮询分配索引，确保负载均衡
    for (size_t i = 0; i < numQueries; ++i) {
        threadIndices[i % numThreads].push_back(i);
    }

    // 启动线程
    for (size_t t = 0; t < numThreads; ++t) {
        threads.emplace_back([&, t] {
            for (size_t idx : threadIndices[t]) {
                double minSq = numeric_limits<double>::max();
                int nearestId = -1;
                nearestSearch(root, queries[idx], minSq, nearestId);
                double distance = (minSq == numeric_limits<double>::max()) ?
                    numeric_limits<double>::quiet_NaN() : sqrt(minSq);
                results[idx] = { nearestId, distance };
            }
            });
    }

    // 等待所有线程完成
    for (auto& thread : threads) {
        thread.join();
    }

    return results;
}